# 1.数据集2014年11月18日至2014年12月18日 user_behavior_20141118_20141218.csv db:user_behavior_20141118_20141218
# 2.数据集2017年11月25日至2017年12月3日 user_behavior_20171125_20171203.csv db:user_behavior_20171125_20171203
# 1.表头 user_id,item_id,category_id,behavior_type,user_geohash,item_category,time
# 2.表头 user_id,item_id,category_id,behavior_type,item_category,time
import gc
import os
import time
from datetime import datetime

import mysql.connector
import pandas as pd
from mysql.connector import Error


def format_time(value):
    try:
        if str(value).replace("-", "").isdigit():
            timestamp = str(value).replace("-", "")
            dt = datetime.fromtimestamp(int(timestamp))
        else:
            dt = datetime.strptime(str(value), '%Y-%m-%d %H')
        return dt.strftime('%Y-%m-%d %H:00:00')
    except Exception as e:
        print(f"时间格式错误：{value}, 错误：{e}")
        return None


def connect_mysql(db_config):
    try:
        connection = mysql.connector.connect(
            host=db_config['host'],
            user=db_config['user'],
            password=db_config['password'],
            database=db_config['database'],
            connection_timeout=28800
        )
        if connection.is_connected():
            print("成功连接到 MySQL 数据库。")
            return connection
    except Error as e:
        print(f"MySQL 连接失败：{e}")
        return None


def insert_batch(cursor, insert_query, batch_data, retries=3):
    attempt = 0
    while attempt < retries:
        try:
            cursor.executemany(insert_query, batch_data)
            return True
        except Error as e:
            print(f"插入失败，重试 {attempt + 1}/{retries} 次。错误：{e}")
            attempt += 1
            time.sleep(2)
    return False


def csv_to_mysql(csv_file, db_config, table_name, batch_size=5000, chunk_size=500000):
    connection = None
    try:
        required_columns = ['user_id', 'item_id', 'item_category', 'behavior_type', 'time']
        connection = connect_mysql(db_config)
        if not connection:
            print("数据库连接失败，程序终止。")
            return

        cursor = connection.cursor()
        columns = ", ".join(required_columns)
        placeholders = ", ".join(["%s"] * len(required_columns))
        insert_query = f"INSERT INTO `{table_name}` ({columns}) VALUES ({placeholders})"

        for chunk in pd.read_csv(csv_file, header=None, encoding='utf-8', chunksize=chunk_size):
            if list(chunk.iloc[0]) == required_columns:
                print("检测到表头，跳过第一行。")
                chunk = chunk[1:]

            if chunk.shape[1] == 6:
                chunk.drop(columns=[2], inplace=True)

            chunk.columns = required_columns
            chunk['time'] = chunk['time'].apply(format_time)
            chunk.dropna(subset=['time'], inplace=True)

            # 显式转换数据类型以减少内存占用
            chunk['user_id'] = chunk['user_id'].astype(int)
            chunk['item_id'] = chunk['item_id'].astype(int)
            chunk['item_category'] = chunk['item_category'].astype(int)

            total_rows = len(chunk)
            for i in range(0, total_rows, batch_size):
                batch_data = chunk.iloc[i:i + batch_size].values.tolist()
                success = insert_batch(cursor, insert_query, batch_data)

                if not success:
                    print(f"尝试重新连接 MySQL 并继续插入数据（从第 {i} 条记录开始）。")
                    connection = connect_mysql(db_config)
                    if not connection:
                        print("多次重连失败，程序终止。")
                        return
                    cursor = connection.cursor()
                    success = insert_batch(cursor, insert_query, batch_data)

                    if not success:
                        print(f"数据从第 {i} 条插入失败，程序终止。")
                        return

                connection.commit()
                print(f"已成功插入 {i + len(batch_data)}/{total_rows} 条数据。")

            # 释放内存
            del chunk
            gc.collect()

        print(f"数据成功插入到表 {table_name}。")

    except Error as e:
        print(f"程序运行错误：{e}")

    finally:
        if connection and connection.is_connected():
            cursor.close()
            connection.close()
            print("MySQL 连接已关闭。")


if __name__ == "__main__":
    csv_file = "user_behavior_20171125_20171203_edit.csv"
    csv_file_path = os.path.join("../static/data", csv_file).replace("\\", "/")
    table_name = csv_file.replace("_edit.csv", "")

    db_config = {
        "host": "localhost",
        "user": "root",
        "password": "123456",
        "database": "user_behavior_db"
    }

    csv_to_mysql(csv_file_path, db_config, table_name)
